#ifndef POINT_PREPROCESS_H
#define POINT_PREPROCESS_H

#include <vector>
#include <cstdint>
#include <iostream>
#include <algorithm>
#include <fstream>
#include <vector>
#include "pointcloud_tools.h"

using namespace pointcloud_tools;

namespace PointPreprocess
{

    struct PointCoord
    {
        int64_t x, y, z;
        int original_index;

        // Generate coord by decoding a 64-bit integer
        PointCoord(uint64_t coord, int idx) : original_index(idx)
        {
            x = coord & 0xFFF;
            y = (coord >> 12) & 0xFFF;
            z = (coord >> 24) & 0xFFF;
        }
    };

    struct OctreeConfig
    {
        int block_size;
        uint8_t *search_range;
        uint16_t *spatial_shape;
        int ich_loop_num;
        int och_loop_num;

        OctreeConfig() : block_size(1024), search_range(nullptr), spatial_shape(nullptr),
                         ich_loop_num(0), och_loop_num(0) {}
    };

    struct PartitionResult
    {
        std::vector<std::vector<int>> pt_mapping;
        std::vector<BlockConfig> block_configs;
        std::vector<std::vector<int>> valid_filters;

        void clear()
        {
            pt_mapping.clear();
            block_configs.clear();
            valid_filters.clear();
        }

        int getBlockCount() const
        {
            return block_configs.size();
        }
    };

    inline bool isInRange(const PointCoord &point,
                          int x_low, int x_high,
                          int y_low, int y_high,
                          int z_low, int z_high)
    {
        return point.x >= x_low && point.x <= x_high &&
               point.y >= y_low && point.y <= y_high &&
               point.z >= z_low && point.z <= z_high;
    }

    void recursive_split(int z_low, int y_low, int x_low,
                         int z_high, int y_high, int x_high,
                         const std::vector<PointCoord> &all_coords,
                         const std::vector<int> &candidate_indices,
                         PartitionResult &result,
                         BlockConfig &bcfg_global,
                         uint64_t &current_pt_addr,
                         uint64_t &current_in_addr,
                         uint64_t &current_rs_addr,
                         const OctreeConfig &config,
                         bool offline_mode = false);

    inline bool partition(BlockConfig &bcfg,
                          const OctreeConfig &config,
                          PartitionResult &result,
                          bool offline_mode = false);

    inline void mergeFeatures(const PartitionResult &partition_result,
                              int valid_pt_num,
                              int och_loop_num,
                              std::vector<uint64_t> &merged_res,
                              const uint64_t (&bias_cores)[4],
                              uint8_t relu_enable = 1);

    void exportResults(const std::vector<uint64_t> &merged_res,
                       int valid_pt_num,
                       int och_loop_num,
                       const std::string &filename);

    void octree_partition(BlockConfig &bcfg,
                          const int block_size,
                          std::vector<std::vector<int>> &pt_mapping,
                          std::vector<BlockConfig> &bcfg_output,
                          std::vector<std::vector<int>> &valid_filter,
                          uint8_t *search_range,
                          uint16_t *spatial_shape,
                          int ich_loop_num,
                          int och_loop_num,
                          bool offline_mode = false);

    void merge_RESULT_feat(const std::vector<BlockConfig> &bcfg_partition,
                           int valid_pt_num,
                           int och_loop_num,
                           const std::vector<std::vector<int>> &pt_mapping,
                           const std::vector<std::vector<int>> &valid_filter,
                           std::vector<uint64_t> &merged_res,
                           const uint64_t (&bias_cores)[4],
                           uint8_t relu_enable = 1);

    void export_merged_results(const std::vector<uint64_t> &merged_res,
                               int valid_pt_num,
                               int och_loop_num,
                               const std::string &filename);

    bool export_blocks_to_files(
        const std::vector<BlockConfig> &block_configs,
        const std::vector<std::vector<int>> &pt_mapping,
        const std::vector<std::vector<int>> &valid_filter,
        const std::string &output_dir,
        const std::vector<uint64_t> &pt_coords,
        const std::vector<uint64_t> &pt_features);

} // namespace PointPreprocess

#endif // POINT_PREPROCESS_H